
/**
 *  算法第四版，281页 -- 红黑树的插入算法
 */
public class RedBlackBST <Key extends Comparable<Key>, Value> {
    private Node root;
    private static final boolean RED = true;
    private static final boolean BLACK = false;

    private class Node {
        // 含有color变量的Node对象
        Key key; // 键
        Value val; // 相关联的值
        Node left, right; // 左右子树
        int N; // 这颗子树中的节点总数
        boolean color; // 由其父节点指向它的链接的颜色

        Node(Key key, Value val, int N, boolean color) {
            this.key = key;
            this.val = val;
            this.N = N;
            this.color = color;
        }
    }

    /**
     *  判断颜色
     * @param x 节点
     * @return 红色返回true
     */
    private boolean isRed(Node x) {
        if(x == null) {
            return false;
        }
        return x.color == RED;
    }

    /**
     *  左旋转h的右链接
     * @param h 旋转的节点
     * @return 返回旋转后位置上的节点
     */
    private Node rotateLeft(Node h) {
        Node x = h.right;
        h.right = x.left;
        x.left = h;
        x.color = h.color;
        h.color = RED;
        x.N = h.N;
        h.N = 1 + size(h.left) + size(h.right);
        return x;
    }

    /**
     *  右旋转h的左链接
     * @param h 旋转的节点
     * @return 返回旋转后位置上的节点
     */
    private Node rotateRight(Node h) {
        Node x = h.left;
        h.left = x.right;
        x.right = h;
        x.color = h.color;
        h.color = RED;
        x.N = h.N;
        h.N = 1 + size(h.left) + size(h.right);
        return x;
    }

    /**
     *  通过转换链接的颜色来分解4-节点
     * @param h
     */
    private void flipColors(Node h) {
        h.color = RED;
        h.left.color = BLACK;
        h.right.color = BLACK;
    }

    /**
     *  返回该节点子树节点的个数
     * @return 返回个数
     */
    private int size() {
        return size(root);
    }
    private int size(Node x) {
        if(x == null) {
            return 0;
        } else {
            return x.N;
        }
    }

    public void put(Key key, Value val) {
        // 查找key，找到则更新其值，否则为它新建一个节点
        root = put(root, key ,val);
        root.color = BLACK;
    }
    private Node put(Node h, Key key, Value val) {
        if(h == null) {
            return new Node(key, val, 1, RED);
        }
        int cmp = key.compareTo(h.key);
        if(cmp < 0) {
            h.left = put(h.left, key, val);
        } else if(cmp > 0) {
            h.right = put(h.right, key, val);
        } else {
            h.val = val;
        }

        if(isRed(h.right) && !isRed(h.left)) {
            h = rotateLeft(h);
        }
        if(isRed(h.left) && isRed(h.left.left)) {
            h = rotateRight(h);
        }
        if(isRed(h.left) && isRed(h.right)) {
            flipColors(h);
        }

        h.N = size(h.left) + size(h.right) + 1;
        return h;
    }

}
